import os
from typing import TypedDict, List
from langchain.schema import HumanMessage, AIMessage


# 定义状态：存储对话消息列表
class ChatState(TypedDict):
    messages: List[HumanMessage | AIMessage]  # 状态字段：对话历史,


from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph

# 初始化 LLM
llm = ChatOpenAI(
    model_name="qwen-plus",
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    temperature=0,
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
)


# 定义节点逻辑：调用 LLM 生成回答，更新状态
def llm_answer_node(state: ChatState) -> ChatState:
    # 1. 从状态中读取对话历史
    messages = state["messages"]
    print(f"对话历史:{messages}")

    # 2. 调用 LLM 生成回答
    response = llm.invoke(messages)
    print(f"大模型返回的答案:{response}")

    # 3. 更新状态：将 AI 回答添加到对话历史
    return {
        "messages": messages + [response]  # 返回状态增量（仅更新 messages 字段）
    }


# 1. 初始化图（指定状态 schema）
graph = StateGraph(ChatState)

# 2. 添加节点（name="llm_answer", 逻辑=llm_answer_node）
graph.add_node("llm_answer", llm_answer_node)

# 3. 定义边：入口→llm_answer 节点（无条件跳转）
graph.set_entry_point("llm_answer")

# 4. 定义出口：llm_answer 节点执行后，流程终止
graph.set_finish_point("llm_answer")

# 5. 编译图（生成可运行的图对象）
compiled_graph = graph.compile()

# 生成 Mermaid 可视化图片
graph_image = compiled_graph.get_graph().draw_mermaid_png()
# 保存图片到文件
with open("graph-01.png", "wb") as f:
    f.write(graph_image)
print("图表已保存为 graph-01.png 文件，请在项目根目录查看")

# 6. 运行图：传入初始状态（用户的第一个问题）
initial_state = {
    "messages": [HumanMessage(content="什么是llm-MCP")]
}

# 执行流程并获取结果
result1 = compiled_graph.invoke(initial_state)

# 打印最终对话
for msg in result1["messages"]:
    print(f"{msg.type.upper()}: {msg.content}\n")

print("======================================")
initial_state2 = {
    "messages": [HumanMessage(content="什么是RAG")]
}

# 执行流程并获取结果
result2 = compiled_graph.invoke(initial_state2)

# 打印最终对话
for msg in result2["messages"]:
    print(f"{msg.type.upper()}: {msg.content}\n")
